/* Write the contents of the <certfile> into kernel symbol system_extra_cert
 *
 * Copyright (C) IBM Corporation, 2015
 *
 * Author: Mehmet Kayaalp <mkayaalp@linux.vnet.ibm.com>
 *
 * This software may be used and distributed according to the terms
 * of the GNU General Public License, incorporated herein by reference.
 *
 * Usage: insert-sys-cert [-s <System.map> -b <vmlinux> -c <certfile> -S <certfile2>
 */

#define _GNU_SOURCE
#include <stdio.h>
#include <ctype.h>
#include <string.h>
#include <limits.h>
#include <stdbool.h>
#include <errno.h>
#include <stdlib.h>
#include <stdarg.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/mman.h>
#include <fcntl.h>
#include <unistd.h>
#include <elf.h>
#include <arpa/inet.h>
#define CERT_SYM  "system_extra_cert"
#define USED_SYM  "system_extra_cert_used"
#define LSIZE_SYM "system_certificate_list_size"

#define FORCE_SIG_SYM "builtin_option_sig"
#define FORCE_CRL_SYM "builtin_option_crl_flags"

#define CERT_SYM_SECOND  "system_extra_cert_second"
#define USED_SYM_SECOND  "system_extra_cert_used_second"

#define SYSTEM_EXTRA_CERT 0
#define RTOS_EXTRA_CERT   1

#define CRL_NOADDR_FORCE 1
#define CRL_NOADDR_PSS_FORCE 2

#define info(format, args...) fprintf(stderr, "INFO:    " format, ## args)
#define warn(format, args...) fprintf(stdout, "WARNING: " format, ## args)
#define  err(format, args...) fprintf(stderr, "ERROR:   " format, ## args)

#ifdef RTOS_TARGET_32BITS
#define CURRENT_ELFCLASS ELFCLASS32
#define Elf_Ehdr	Elf32_Ehdr
#define Elf_Shdr	Elf32_Shdr
#define Elf_Sym		Elf32_Sym
#define USE32		true
#else
#define CURRENT_ELFCLASS ELFCLASS64
#define Elf_Ehdr	Elf64_Ehdr
#define Elf_Shdr	Elf64_Shdr
#define Elf_Sym		Elf64_Sym
#define USE32		false
#endif

static u_int64_t convert_to_le_u64(u_int64_t val, bool convert)
{
	if (convert) {
		if (!USE32)
			return (((u_int64_t)ntohl((u_int32_t)((val << 32) >> 32))) << 32) |
				(u_int32_t)ntohl((u_int32_t)(val >> 32));
		else
			return ntohl(val);
	}

	return val;
}

static unsigned convert_to_le_u32(unsigned val, bool convert)
{
	if (convert)
		return ntohl(val);

	return val;
}

static unsigned short convert_to_le_u16(unsigned short val, bool convert)
{
	if (convert)
		return ntohs(val);

	return val;
}

static u_int64_t convert_to_be_u64(u_int64_t val, bool convert)
{
	if (convert) {
		if (!USE32)
			return (u_int32_t)htonl((u_int32_t)(val >> 32)) |
				(((u_int64_t)htonl((u_int32_t)((val << 32) >> 32))) << 32);
		else
			return htonl(val);
	}

	return val;
}

static unsigned char endianness(void)
{
	uint16_t two_byte = 0x00FF;
	uint8_t low_address = *((uint8_t *)&two_byte);

	if (low_address == 0)
		return ELFDATA2MSB;
	else
		return ELFDATA2LSB;
}

struct sym {
	char *name;
	unsigned long address;
	unsigned long offset;
	void *content;
	size_t size;
};

static unsigned long get_offset_from_address(Elf_Ehdr *hdr, unsigned long addr)
{
	Elf_Shdr *x;
	unsigned int i, num_sections;

	x = (void *)hdr + hdr->e_shoff;
	if (hdr->e_shnum == SHN_UNDEF)
		num_sections = x[0].sh_size;
	else
		num_sections = hdr->e_shnum;

	for (i = 1; i < num_sections; i++) {
		unsigned long start = x[i].sh_addr;
		unsigned long end = start + x[i].sh_size;
		unsigned long offset = x[i].sh_offset;

		if (addr >= start && addr <= end)
			return addr - start + offset;
	}
	return 0;
}


#define LINE_SIZE 100

static void get_symbol_from_map(Elf_Ehdr *hdr, FILE *f, char *name,
				struct sym *s)
{
	char l[LINE_SIZE];
	char *w, *p, *n = NULL;

	s->size = 0;
	s->address = 0;
	s->offset = 0;
	if (fseek(f, 0, SEEK_SET) != 0) {
		perror("File seek failed");
		return;
	}
	while (fgets(l, LINE_SIZE, f)) {
		p = strchr(l, '\n');
		if (!p) {
			err("Missing line ending.\n");
			return;
		}
		n = strstr(l, name);
		if (n) {
			/* n should be larger than name with exactly a '\n'. */
			if (n + strlen(name) == p)
				break;
		}
	}
	if (!n) {
		err("Unable to find symbol: %s\n", name);
		return;
	}
	w = strchr(l, ' ');
	if (!w)
		return;

	*w = '\0';
	s->address = strtoul(l, NULL, 16);
	if (s->address == 0)
		return;
	s->offset = get_offset_from_address(hdr, s->address);
	s->name = name;
	s->content = (void *)hdr + s->offset;
}

static Elf_Sym *find_elf_symbol(Elf_Ehdr *hdr, Elf_Shdr *symtab, char *name, bool convert)
{
	Elf_Sym *sym, *symtab_start;
	char *strtab, *symname;
	unsigned int link;
	Elf_Shdr *x;
	size_t i, n;

	x = (void *)hdr + convert_to_le_u64(hdr->e_shoff, convert);
	link = convert_to_le_u32(symtab->sh_link, convert);
	symtab_start = (void *)hdr + convert_to_le_u64(symtab->sh_offset, convert);
	n = convert_to_le_u64(symtab->sh_size, convert) / convert_to_le_u64(symtab->sh_entsize, convert);
	strtab = (void *)hdr + convert_to_le_u64(x[link].sh_offset, convert);

	for (i = 0; i < n; i++) {
		sym = &symtab_start[i];
		symname = strtab + convert_to_le_u32(sym->st_name, convert);
		if (strcmp(symname, name) == 0)
			return sym;
	}
	err("Unable to find symbol: %s\n", name);
	return NULL;
}

static void get_symbol_from_table(Elf_Ehdr *hdr, Elf_Shdr *symtab,
				  char *name, struct sym *s, bool convert)
{
	Elf_Shdr *sec;
	int secndx;
	Elf_Sym *elf_sym;
	Elf_Shdr *x;

	x = (void *)hdr + convert_to_le_u64(hdr->e_shoff, convert);
	s->size = 0;
	s->address = 0;
	s->offset = 0;
	elf_sym = find_elf_symbol(hdr, symtab, name, convert);
	if (!elf_sym)
		return;
	secndx = convert_to_le_u16(elf_sym->st_shndx, convert);
	if (!secndx)
		return;
	sec = &x[secndx];
	s->size = convert_to_le_u64(elf_sym->st_size, convert);
	s->address = convert_to_le_u64(elf_sym->st_value, convert);
	s->offset = s->address - convert_to_le_u64(sec->sh_addr, convert)
			       + convert_to_le_u64(sec->sh_offset, convert);
	s->name = name;
	s->content = (void *)hdr + s->offset;
}

static Elf_Shdr *get_symbol_table(Elf_Ehdr *hdr, bool convert)
{
	Elf_Shdr *x;
	unsigned int i, num_sections;

	x = (void *)hdr + convert_to_le_u64(hdr->e_shoff, convert);
	if (convert_to_le_u16(hdr->e_shnum, convert) == SHN_UNDEF)
		num_sections = convert_to_le_u16(x[0].sh_size, convert);
	else
		num_sections = convert_to_le_u16(hdr->e_shnum, convert);

	for (i = 1; i < num_sections; i++) {
		if (convert_to_le_u32(x[i].sh_type, convert) == SHT_SYMTAB)
			return &x[i];
	}
	return NULL;
}

static void *map_file(char *file_name, size_t *size)
{
	struct stat st;
	void *map;
	int fd;

	fd = open(file_name, O_RDWR);
	if (fd < 0) {
		perror(file_name);
		return NULL;
	}
	if (fstat(fd, &st)) {
		perror("Could not determine file size");
		close(fd);
		return NULL;
	}
	*size = st.st_size;
	map = mmap(NULL, *size, PROT_READ|PROT_WRITE, MAP_SHARED, fd, 0);
	if (map == MAP_FAILED) {
		perror("Mapping to memory failed");
		close(fd);
		return NULL;
	}
	close(fd);
	return map;
}


static char *read_file(char *file_name, size_t *size)
{
	struct stat st;
	char *buf;
	int fd;

	fd = open(file_name, O_RDONLY);
	if (fd < 0) {
		perror(file_name);
		return NULL;
	}
	if (fstat(fd, &st)) {
		perror("Could not determine file size");
		close(fd);
		return NULL;
	}
	*size = st.st_size;
	buf = malloc(*size);
	if (!buf) {
		perror("Allocating memory failed");
		close(fd);
		return NULL;
	}
	if (read(fd, buf, *size) != *size) {
		perror("File read failed");
		close(fd);
		free(buf);
		return NULL;
	}
	close(fd);
	return buf;
}

static void print_sym(Elf_Ehdr *hdr, struct sym *s)
{
	info("sym:    %s\n", s->name);
	info("addr:   0x%lx\n", s->address);
	info("size:   %zu\n", s->size);
	info("offset: 0x%lx\n", (unsigned long)s->offset);
}

static void print_usage(char *e)
{
	printf("Usage %s [-s <System.map>] -b <vmlinux> -c <certfile> -S <certfile2>\n", e);
}

static void insert_cert_from_map(char *cert_file, Elf_Ehdr *hdr, FILE *system_map, int type)
{
	struct sym cert_sym, used_sym;
	size_t cert_size = 0;
	char *cert = NULL;

	cert = read_file(cert_file, &cert_size);
	if (!cert)
		exit(EXIT_FAILURE);

	if (type == SYSTEM_EXTRA_CERT) {
		get_symbol_from_map(hdr, system_map, CERT_SYM, &cert_sym);
		get_symbol_from_map(hdr, system_map, USED_SYM, &used_sym);
	} else {
		get_symbol_from_map(hdr, system_map, CERT_SYM_SECOND, &cert_sym);
		get_symbol_from_map(hdr, system_map, USED_SYM_SECOND, &used_sym);
	}

	info("Symbol found from map.\n");
	if (cert_sym.offset == 0 || used_sym.offset == 0) {
		err("No offset info found from map!\n");
		exit(EXIT_FAILURE);
	}

	if (used_sym.address < cert_sym.address) {
		err("used_sym.address is less than cert_sym.address!\n");
		exit(EXIT_FAILURE);
	}

	cert_sym.size = used_sym.address - cert_sym.address;
	if (cert_sym.size < cert_size) {
		err("Certificate is larger than the reserved area!\n");
		exit(EXIT_FAILURE);
	}

	memcpy(cert_sym.content, cert, cert_size);
	if (cert_size < cert_sym.size)
		memset(cert_sym.content + cert_size, 0, cert_sym.size - cert_size);
}

/*
 * enable_crl_noaddr -- enable force crl verification in kernel
 *
 * Return: true if succeed, false if failed
 */
static bool enable_crl_noaddr(Elf_Ehdr *hdr, Elf_Shdr *symtab, bool convert)
{
	struct sym force_crl_sym = { 0 };
	unsigned long *crl_noaddr = NULL;

	get_symbol_from_table(hdr, symtab, FORCE_CRL_SYM, &force_crl_sym, convert);
	if (force_crl_sym.offset == 0 || force_crl_sym.content == NULL) {
		err("No 'crl_noaddr' symbol info found!\n");
		return false;
	}

	crl_noaddr = (unsigned long *)force_crl_sym.content;
	info("%s is %lu\n", FORCE_CRL_SYM, *crl_noaddr);
	*crl_noaddr = convert_to_be_u64(CRL_NOADDR_FORCE, convert);
	info("%s is updated to %lu\n", FORCE_CRL_SYM, *crl_noaddr);

	return true;
}

/*
 * enable_crl_noaddr_pss -- enable force crl_pss verification in kernel
 *
 * Return: true if succeed, false if failed
 */
static bool enable_crl_noaddr_pss(Elf_Ehdr *hdr, Elf_Shdr *symtab, bool convert)
{
	struct sym force_crl_pss_sym = { 0 };
	unsigned long *crl_noaddr_pss = NULL;

	get_symbol_from_table(hdr, symtab, FORCE_CRL_SYM, &force_crl_pss_sym, convert);
	if (force_crl_pss_sym.offset == 0 || force_crl_pss_sym.content == NULL) {
		err("No 'crl_noaddr_pss' symbol info found!\n");
		return false;
	}

	crl_noaddr_pss = (unsigned long *)force_crl_pss_sym.content;
	info("%s is %lu\n", FORCE_CRL_SYM, *crl_noaddr_pss);
	*crl_noaddr_pss = convert_to_be_u64(CRL_NOADDR_PSS_FORCE, convert);
	info("%s is updated to %lu\n", FORCE_CRL_SYM, *crl_noaddr_pss);

	return true;
}

/*
 * enable_sig_enforce -- enable sig_enforce in kernel
 *
 * Return: true if succeed, false if failed
 */
static bool enable_sig_enforce(Elf_Ehdr *hdr, Elf_Shdr *symtab, bool convert)
{
	struct sym force_sig_sym = { 0 };
	unsigned long *builtin_option_sig = NULL;

	get_symbol_from_table(hdr, symtab, FORCE_SIG_SYM, &force_sig_sym, convert);
	if (force_sig_sym.offset == 0 || force_sig_sym.content == NULL) {
		err("No 'sig_enforce' symbol info found!\n");
		return false;
	}

	builtin_option_sig = (unsigned long *)force_sig_sym.content;
	info("%s is %lu\n", FORCE_SIG_SYM, *builtin_option_sig);
	*builtin_option_sig = convert_to_be_u64(1UL, convert);
	info("%s is updated to %lu\n", FORCE_SIG_SYM, *builtin_option_sig);

	return true;
}

static void update_lsize(const struct sym *used_sym, const struct sym *lsize_sym,
			 size_t cert_size, bool convert)
{
	unsigned long *used = (unsigned long *)used_sym->content;
	unsigned long *lsize = (unsigned long *)lsize_sym->content;

	info("%s is %lu\n", LSIZE_SYM, *lsize);
	*lsize = convert_to_be_u64(*lsize + cert_size - *used, convert);
	info("%s is update to %lu\n", LSIZE_SYM, *lsize);
}

static size_t copy_cert(const struct sym *cert_sym, const struct sym *used_sym,
						const char *cert, size_t cert_size,
						bool convert)
{
	unsigned long used_host;
	unsigned long *used;

	used = (unsigned long *)used_sym->content;
	used_host = convert_to_le_u64(*used, convert);

	if (cert_sym->size < cert_size) {
		err("Certificate is larger than the reserved area!\n");
		exit(EXIT_FAILURE);
	}

	/* If the existing cert is the same, don't overwrite */
	if (cert_size == used_host && memcmp(cert_sym->content, cert, cert_size) == 0) {
		warn("Certificate was already inserted.\n");
		exit(EXIT_SUCCESS);
	}

	if (used_host > 0)
		warn("Replacing previously inserted certificate.\n");

	memcpy(cert_sym->content, cert, cert_size);
	if (cert_size < cert_sym->size)
		memset(cert_sym->content + cert_size, 0, cert_sym->size - cert_size);

	*used = convert_to_be_u64(cert_size, convert);
	return convert_to_le_u64(*used, convert);
}

static void insert_cert_from_table(char *cert_file, Elf_Ehdr *hdr, bool convert, int type)
{
	struct sym cert_sym, used_sym, lsize_sym;
	Elf_Shdr *symtab = NULL;
	size_t cert_size = 0;
	char *cert = NULL;
	unsigned long used_host;

	cert = read_file(cert_file, &cert_size);
	if (!cert) {
		err("No cert is found!\n");
		exit(EXIT_FAILURE);
	}

	symtab = get_symbol_table(hdr, convert);
	if (!symtab) {
		err("No symbol table found.\n");
		exit(EXIT_FAILURE);
	}
	info("Symbol table found.\n");
	if (type == SYSTEM_EXTRA_CERT) {
		get_symbol_from_table(hdr, symtab, CERT_SYM, &cert_sym, convert);
		get_symbol_from_table(hdr, symtab, USED_SYM, &used_sym, convert);
		get_symbol_from_table(hdr, symtab, LSIZE_SYM, &lsize_sym, convert);
	} else {
		get_symbol_from_table(hdr, symtab, CERT_SYM_SECOND, &cert_sym, convert);
		get_symbol_from_table(hdr, symtab, USED_SYM_SECOND, &used_sym, convert);
	}
	info("Symbol found.\n");
	if (cert_sym.offset == 0 || used_sym.offset == 0) {
		err("No offset info found!\n");
		exit(EXIT_FAILURE);
	}
	if (type == SYSTEM_EXTRA_CERT && lsize_sym.offset == 0) {
		err("No lsize offset info found!\n");
		exit(EXIT_FAILURE);
	}

	print_sym(hdr, &cert_sym);
	print_sym(hdr, &used_sym);

	if (type == SYSTEM_EXTRA_CERT) {
		print_sym(hdr, &lsize_sym);
		update_lsize(&used_sym, &lsize_sym, cert_size, convert);
	}

	/* if copy_cert succeed, it will return the sizeof cert written. */
	used_host = copy_cert(&cert_sym, &used_sym, cert, cert_size, convert);

	info("Inserted the contents of %s into %lx.\n", cert_file, cert_sym.address);
	info("Used %lu bytes out of %zu bytes reserved.\n", used_host, cert_sym.size);
}

int main(int argc, char **argv)
{
	char *system_map_file = NULL;
	char *vmlinux_file = NULL;
	char *cert_file = NULL;
	char *cert_file_second = NULL;
	bool builtin_option_sig = false;
	/* When it equals 1, it's crl_noaddr.When it equals 2, it's crl_noaddr_pss. */
	long crl_noaddr_args = 0;
	char *ptr;

	size_t vmlinux_size = 0;
	FILE *system_map;
	Elf_Ehdr *hdr;
	int opt;
	Elf_Shdr *symtab = NULL;
	bool convert = false;

	while ((opt = getopt(argc, argv, "b:c:s:S:er:")) != -1) {
		switch (opt) {
		case 's':
			system_map_file = optarg;
			break;
		case 'b':
			vmlinux_file = optarg;
			break;
		case 'c':
			cert_file = optarg;
			break;
		case 'S':
			cert_file_second = optarg;
			break;
		case 'e':
			builtin_option_sig = true;
			info("builtin_option_sig is set to true.\n");
			break;
		case 'r':
			errno = 0;
			crl_noaddr_args = strtol(optarg, &ptr, 10);
			if (errno == ERANGE) {
				err("strtol error\n");
				exit(EXIT_FAILURE);
			}
			info("crl_noaddr_args is set to : %ld\n", crl_noaddr_args);
			break;
		default:
			break;
		}
	}

	if (!vmlinux_file || !cert_file) {
		print_usage(argv[0]);
		exit(EXIT_FAILURE);
	}

	hdr = map_file(vmlinux_file, &vmlinux_size);
	if (!hdr)
		exit(EXIT_FAILURE);

	if (vmlinux_size < sizeof(*hdr)) {
		err("Invalid ELF file.\n");
		exit(EXIT_FAILURE);
	}

	if ((hdr->e_ident[EI_MAG0] != ELFMAG0) ||
	    (hdr->e_ident[EI_MAG1] != ELFMAG1) ||
	    (hdr->e_ident[EI_MAG2] != ELFMAG2) ||
	    (hdr->e_ident[EI_MAG3] != ELFMAG3)) {
		err("Invalid ELF magic.\n");
		exit(EXIT_FAILURE);
	}

	if (hdr->e_ident[EI_CLASS] != CURRENT_ELFCLASS) {
		err("ELF class mismatch.\n");
		exit(EXIT_FAILURE);
	}

	if (hdr->e_ident[EI_DATA] != endianness()) {
		info("ELF endian mismatch, will convert to little endian.\n");
		convert = true;
	}

	if (convert_to_le_u64(hdr->e_shoff, convert) > vmlinux_size) {
		err("Could not find section header.\n");
		exit(EXIT_FAILURE);
	}

	symtab = get_symbol_table(hdr, convert);
	if (!symtab) {
		warn("Could not find the symbol table.\n");
		if (!system_map_file) {
			err("Please provide a System.map file.\n");
			print_usage(argv[0]);
			exit(EXIT_FAILURE);
		}

		system_map = fopen(system_map_file, "r");
		if (!system_map) {
			perror(system_map_file);
			exit(EXIT_FAILURE);
		}
		if (cert_file)
			insert_cert_from_map(cert_file, hdr, system_map, SYSTEM_EXTRA_CERT);

		if (cert_file_second)
			insert_cert_from_map(cert_file_second, hdr, system_map, RTOS_EXTRA_CERT);

		fclose(system_map);
	} else {
		if (cert_file)
			insert_cert_from_table(cert_file, hdr, convert, SYSTEM_EXTRA_CERT);

		if (cert_file_second)
			insert_cert_from_table(cert_file_second, hdr, convert, RTOS_EXTRA_CERT);

		if (crl_noaddr_args == CRL_NOADDR_FORCE && !enable_crl_noaddr(hdr, symtab, convert)) {
			err("crl_noaddr (-m) is given, but failed to enable force crl\n");
			exit(EXIT_FAILURE);
		}

		if (builtin_option_sig && !enable_sig_enforce(hdr, symtab, convert)) {
			err("sig_enforce (-e) is given, but failed to enable force signature\n");
			exit(EXIT_FAILURE);
		}

		if (crl_noaddr_args == CRL_NOADDR_PSS_FORCE && !enable_crl_noaddr_pss(hdr, symtab, convert)) {
			err("crl_noaddr_pss (-p) is given, but failed to enable force crl_pss\n");
			exit(EXIT_FAILURE);
		}
	}
	exit(EXIT_SUCCESS);
}
